from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from config import Config

Bcenter_DB_URL =  f"mysql+pymysql://{Config.DB_USER}:{Config.DB_PASS}@{Config.DB_HOST}:{Config.DB_PORT}/{Config.DB_NAME}"

engine = create_engine(
    Bcenter_DB_URL,
    pool_pre_ping=True,  # 每次使用连接前检查连接是否有效
    pool_recycle=300,  # 每5分钟回收连接（避免超过MySQL的wait_timeout）
    pool_size=2,
    max_overflow=20,
    connect_args={
        "connect_timeout": 10  # 连接超时设置
    }
)

def get_db(DB_KEY: str = "Bcenter"):
    def _get_session():
        DB_URLS ={
            "Bcenter":    Bcenter_DB_URL,
            "Omscenter": "mysql+pymysql://cread:jtw6huC6skmWWB99@172.16.51.8:6001/test",
            "Listingcenter": "mysql+pymysql://cread:jtw6huC6skmWWB99@172.16.51.8:6002/test",
            "Icenter": "mysql+pymysql://cread:jtw6huC6skmWWB99@172.16.51.8:6003/test",
            "Pcenter": "mysql+pymysql://cread:jtw6huC6skmWWB99@172.16.51.8:6004/db_replenishment_sales_center",
            "Gcenter": "mysql+pymysql://cread:jtw6huC6skmWWB99@172.16.51.8:6005/test",
            }
        engine = create_engine(
            DB_URLS[DB_KEY],
            pool_pre_ping=True,  # 每次使用连接前检查连接是否有效
            pool_recycle=300,  # 每5分钟回收连接（避免超过MySQL的wait_timeout）
            pool_size=2,
            max_overflow=20,
            connect_args={
                "connect_timeout": 10  # 连接超时设置
            }
        )
        session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
        return session()

    db = _get_session()
    try:
        yield db
    finally:
        db.close()


